from datetime import datetime

import streamlit as st
import requests
from typing import List, Dict
from elasticsearch import Elasticsearch
from langchain_community.document_compressors import JinaRerank
from langchain_community.embeddings import JinaEmbeddings
from openai import OpenAI

from es_retriever import get_relevant_documents

user_id = 10086
# 配置DeepSeek API
client = OpenAI(
    base_url='https://api-inference.modelscope.cn/v1/',
    api_key='b011f8fd-99a5-4633-8bf5-fef46cf7bd90',  # ModelScope Token
)

# 自定义CSS样式
STYLES = """
<style>
.think-block {
    background-color: #f0f0f0;
    border-left: 4px solid #4CAF50;
    padding: 1rem;
    margin: 1rem 0;
    font-style: italic;
    color: #555;
    border-radius: 0 5px 5px 0;
}
.container-border {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 1rem;
    margin-bottom: 1rem;
}
</style>
"""
es_url = "http://10.252.202.55:9200"  # Elasticsearch地址
es = Elasticsearch([es_url])
rag_index = "rag_documents"
file_index = "rag_file_info"
# 存储到Elasticsearch
embeddings = JinaEmbeddings(
    jina_api_key="jina_21eda124ed7e4ee895f7bb7197f3cc5ebT45ebFpQaWsIeieL6QYTY9Sojgu",
    model_name="jina-embeddings-v2-base-zh"
)

compressor = JinaRerank(jina_api_key="jina_21eda124ed7e4ee895f7bb7197f3cc5ebT45ebFpQaWsIeieL6QYTY9Sojgu",
                        model="jina-reranker-v2-base-multilingual")

es_url = "http://10.252.202.55:9200"  # Elasticsearch地址
es_client = Elasticsearch(hosts=[es_url])
print(es_client.info())


class DeepSeekChatApp:
    def __init__(self):
        # 初始化会话状态
        if 'messages' not in st.session_state:
            st.session_state.messages = []

    def build_system_prompt(self, final_docs: List, query: str) -> str:
        """构建系统提示模板"""
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        search_results = ""
        for i, item in enumerate(final_docs, start=1):
            chunk_id = item.metadata['chunk_id']
            filename = item.metadata['filename']
            page = item.page_content[:2000]
            search_results += (
                f"[chunk {chunk_id} begin]\n"
                f"filename: {filename}\n"
                f"content: {page}...\n"  # 限制内容长度
                f"[chunk {chunk_id} end]\n"
            )

        return f"""
             # 以下内容是基于用户发送的消息的搜索结果:{search_results}

             在我给你的搜索结果中，每个结果都是[chunk X begin]...[chunk X end]格式的，X代表每篇文章某个切片的数字索引，一个文章可能包含多个切片，顺序是乱序。
             请在适当的情况下在句子末尾引用上下文。请按照引用编号[citation:X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[citation:3][citation:5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
             在回答时，请注意以下几点：
             - 今天是{current_time}。
             - 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
             - 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
             - 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[citation:3][citation:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
             - 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
             - 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
             - 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
             - 你的回答应该综合多个文件来回答，不能重复引用一个文件。
             - 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
             # 用户消息为：{query}"""

    def query_rag(self, query):
        # 混合检索
        docs = get_relevant_documents(query, es_client=es_client,
                                      index_name=rag_index,
                                      embedding_model=embeddings,
                                      keyword_boost=0.4,
                                      vector_boost=0.6,
                                      default_filters=[
                                          {"term": {"metadata.user_id": user_id}}  # 默认过滤条件
                                      ])
        # 重排序
        re_docs = compressor.rerank(query=query, documents=docs, top_n=3)
        print(re_docs)
        final_docs = []
        for doc in re_docs:
            index = doc['index']
            final_docs.append(docs[index])
        return self.build_system_prompt(final_docs, query)

    def get_deepseek_response(self, messages: List[Dict], temperature: float = 0.7) -> str:
        """调用DeepSeek API获取响应"""
        try:
            response = client.chat.completions.create(
                model='deepseek-ai/DeepSeek-R1-Distill-Qwen-7B',  # ModelScope Model-Id
                messages=messages,
                stream=True
            )
            done_reasoning = False
            for chunk in response:
                reasoning_chunk = chunk.choices[0].delta.reasoning_content
                answer_chunk = chunk.choices[0].delta.content
                if reasoning_chunk != '':
                    print(reasoning_chunk, end='', flush=True)
                    yield 'think', reasoning_chunk
                elif answer_chunk != '':
                    if not done_reasoning:
                        print('\n\n === Final Answer ===\n')
                        done_reasoning = True
                    print(answer_chunk, end='', flush=True)
                    yield 'content', answer_chunk

        except requests.RequestException as e:
            st.error(f"API调用错误: {str(e)}")
            yield f"API调用错误: {str(e)}"

    def render_chat_interface(self):
        """渲染聊天界面"""
        st.markdown(STYLES, unsafe_allow_html=True)
        st.title("🤖 知识库小助手")

        # 侧边栏配置
        with st.sidebar:
            st.header("对话设置")
            temperature = st.slider("创造力", 0.0, 1.0, 0.7, step=0.1)
            system_prompt = st.text_area("系统角色设定", "你是一个友好、专业的AI助手")

        # 显示历史消息
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # 用户输入
        if prompt := st.chat_input("输入您的问题"):
            # 添加用户消息
            st.session_state.messages.append({"role": "user", "content": prompt})

            # 显示用户消息
            with st.chat_message("user"):
                st.markdown(prompt)

            # TODO 加一个意图判断

            # 检索
            system_prompt = self.query_rag(prompt)

            # 准备API消息列表
            messages = [
                {"role": "system", "content": system_prompt},
                *[{"role": m["role"], "content": m["content"]} for m in st.session_state.messages]
            ]

            # 显示AI响应
            with st.chat_message("assistant"):
                response_placeholder = st.empty()
                full_response = ""
                full_think_response = ""
                think_count = 0
                # 流式显示响应
                for tag, chunk in self.get_deepseek_response(messages, temperature):
                    if 'think' == tag:
                        if think_count == 0:
                            full_think_response += f'<div class="think-block">'
                        think_count = 1
                        full_think_response += chunk
                        response_placeholder.markdown(full_think_response + "▌", unsafe_allow_html=True)
                    else:
                        if think_count == 1:
                            full_think_response += '</div>'
                            response_placeholder.markdown(full_think_response, unsafe_allow_html=True)
                            think_count = 0
                        full_response += chunk
                        response_placeholder.markdown(full_think_response + full_response + "▌", unsafe_allow_html=True)

                # response_placeholder.markdown(full_think_response + full_response)

            # 保存AI消息
            st.session_state.messages.append({"role": "assistant", "content": full_response})


def main():
    chat_app = DeepSeekChatApp()
    chat_app.render_chat_interface()


if __name__ == "__main__":
    main()
